from add_nms_layer2 import add_topk_postprocess
from add_onnx_infer import infer_by_onnx
from check_add_layer import test_model_by_onnxruntime

# 参数配置
num_classes = 3 # 类别数量
num_keypoints = 6 # 关键点数量
k = 10  # 设置 Top-K 数量 k 表示取得前 k 个最大置信度检测框
# confidence_threshold = 0.5
# nms_threshold = 0.45
# max_output_boxes_per_class = 9
had_visible = True

# ONNX 模型路径
origin_model_path = r'/home/champrin/Desktop/rm24_arm_ws/src/arm_auto_exchange/exchange_slot_detector/model/exchange_best421.onnx'
output_model_path = r'../exchange_best421_v8_pose_nms_concat_topk.onnx'

# 图像路径
image_path = r'/home/champrin/Desktop/MV-CS016-10UC+DA1041860/a/exchange_red_Video_20250113103714500_146.jpg'

# 对原始 YOLOv8 / 10 模型添加 topK 层，使得推理时就进行大部分后处理
# 原始输出为 4（矩形框的中心点x y与width height） + num_classes（num_classes个类别置信度） + num_keypoints * 3 （num_keypoints个 关键点坐标xy 与 是否可见）
# 修改后输出维度为 batch x k x (4 + 1 + 1 + num_keypoints * 3)
# 4 + 1 + 1 + num_keypoints * 3 代表 4（矩形框的中心点x y与width height） + 1（最大类别的置信度） + 1（最大类别id） + num_keypoints * 3 （num_keypoints个 关键点坐标xy 与 关键点置信度）
add_topk_postprocess(origin_model_path, num_classes, num_keypoints, k, output_model_path, had_visible)

# 打印 onnx 模型各层的输出维度，看是否合乎设计
ort_outs = test_model_by_onnxruntime(output_model_path)

for key, value in ort_outs.items():
    print(f"Key: {key}, Value: {value.shape}")

# 使用修改过的添加 topK 层进行推理验证
infer_by_onnx(output_model_path, image_path)

